//      vom_builder.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include "vom_builder.h"

void vom::node::set_frequency(const vom::sastring& train)
{
  std::string s = this->str();
  int n = sz(s);

  int left, size;
  train.sasearch(s, left, size);
	
  BOOST_FOREACH(FrequencyTable& ft, f_) ft.zeros();

  n_occurence_ = 0;
  for (int i = 0; i < size; i++) { // all occurences of context s
    int pos = train.SA[left+i];
    if (pos + n < sz(train)) { 
      char c = train[pos+n];
      for (int j = 0; j < (int)strlen(alphabet_); j++)
	if (alphabet_[j] == c) f_[(pos+n) % periodic()][j] += 1.0;
    }
  }
  n_occurence_ += size;
  normalize();
}

void vom::node::set_frequency(const std::vector<vom::sastring>& train_set)
{
  std::string s = this->str();
  int n = sz(s);

  BOOST_FOREACH(FrequencyTable& ft, f_) ft.zeros();

  n_occurence_ = 0;
  BOOST_FOREACH(const vom::sastring& train, train_set) {
    int left, size;
    train.sasearch(s, left, size);
		
    for (int i = 0; i < size; i++) { // all occurences of context s
      int pos = train.SA[left+i];
      if (pos + n < sz(train)) { 
	char c = train[pos+n];
	for (int j = 0; j < (int)strlen(alphabet_); j++)
	  if (alphabet_[j] == c) f_[(pos+n) % periodic()][j] += 1.0;
      }
    }
    n_occurence_ += size;
  }
  normalize();
}

void vom::node::set_rand_frequency()
{
  BOOST_FOREACH(FrequencyTable& ft, f_) ft.randu();
  normalize();
}

void vom::node::normalize()
{
  BOOST_FOREACH(FrequencyTable& ft, f_) {
    double min = DOUBLE_INFINITY;
    
    if (arma::sum(ft) < 2) {
      ft = arma::vec(); // delete this table as estimate is not reliable
      continue;
    }

    for (int i = 0; i < sz(ft); i++)
      if (min > ft[i] && ft[i] > 1e-16) min = ft[i];
    if (min == DOUBLE_INFINITY) min = 1;
    
    ft += min * 0.1;
    double sum = arma::sum(ft);
    ft /= sum;
  }
}

std::string vom::node::str() const
{
  const node* N = this;
  std::string s;
	
  while (N->parent_ != NULL) {
    s = s + N->c_;
    N = N->parent_;
  }
  return s;
}

void vom::VOM::build_fom(const vom::sastring& train, int order)
{
  root_ = new node('X', alphabet.c_str(), NULL, periodic_);
  build_fom_node(root_, train, order);
}

void vom::VOM::build_fom_node(node* N, const vom::sastring& train, int order)
{
  N->set_frequency(train);
  if (order > 0) {
    for (int i = 0; i < sz(alphabet); i++) {
      char c = alphabet[i];
      node* child = new node(c, alphabet.c_str(), N, periodic_);
      N->child_[i] = child;
    }
    BOOST_FOREACH(node* child, N->child_)
      build_fom_node(child, train, order-1);
  }
}

void vom::VOM::print_node(std::ostream& os, node* N, int depth)
{
  if (N == NULL) {
    os << "UP" << ENDL;
    return;
  }
  //for (int i = 0; i < depth; i++) os << "  ";

  if (N->parent_ == NULL) os << "X";
  else os << N->str();
  BOOST_FOREACH(FrequencyTable& ft, N->f_) {
    if (sz(ft) == sz(alphabet))
      for (int i = 0; i < sz(ft); i++) os << " " << ft[i];
    else
      for (int i = 0; i < sz(alphabet); i++) os << " " << 0;
  }
  os << ENDL;
	
  BOOST_FOREACH(node* child, N->child_)
    print_node(os, child, depth+1);
}

void vom::VOM::load_node(std::istream& is, node** N, node* parent, int depth)
{
  FileAccessLineSplit(is, line, tokens, "\r\n\t ,");
  if (sz(tokens) == 0 || tokens[0] == "UP") {
    *N = NULL;
    return;
  }
  count_ ++;
	
  (*N) = new node(tokens[0][0], alphabet.c_str(), parent, periodic_);
  for (int i = 0; i < sz(alphabet)*periodic_; i++) {
    double f = RAND::s2type<double>(tokens[i+1]);
    int frame = i/sz(alphabet);
    (*N)->f_[frame][i%sz(alphabet)] = f;
  }
  for (int fr = 0; fr < periodic_; fr++)
    if (arma::sum((*N)->f_[fr]) == 0) (*N)->f_[fr] = arma::vec();
	
  for (int i = 0; i < sz(alphabet); i++) {
    node** child = &((*N)->child_[i]);
    load_node(is, child, *N, depth+1);
  }
}


void vom::VOM::build_fom(const std::vector<vom::sastring>& train_set, int order)
{
  root_ = new node('X', alphabet.c_str(), NULL, periodic_);
  build_fom_node(root_, train_set, order);
}

void vom::VOM::build_rand_fom(int order)
{
  root_ = new node('X', alphabet.c_str(), NULL, periodic_);
  build_rand_fom_node(root_, order);
}

void vom::VOM::build_fom_node(node* N, const std::vector<vom::sastring>& train_set, int order)
{
  N->set_frequency(train_set);
  if (order > 0) {
    for (int i = 0; i < sz(alphabet); i++) {
      char c = alphabet[i];
      node* child = new node(c, alphabet.c_str(), N, periodic_);
      N->child_[i] = child;
    }
    BOOST_FOREACH(node* child, N->child_)
      build_fom_node(child, train_set, order-1);
  }
}

void vom::VOM::build_rand_fom_node(node* N, int order)
{
  N->set_rand_frequency();
  if (order > 0) {
    for (int i = 0; i < sz(alphabet); i++) {
      char c = alphabet[i];
      node* child = new node(c, alphabet.c_str(), N, periodic_);
      N->child_[i] = child;
    }
    BOOST_FOREACH(node* child, N->child_)
      build_rand_fom_node(child, order-1);
  }
}

double vom::VOM::log_prob(const std::string& s, int pos)
{
  node* N = root_;
  int i = pos-1;
	
  bool stop = false;
  int frame = pos % periodic_;
  while (i >= 0 && !stop) { // trace upstream for longest context
    char c = s[i];
    stop = true;
    int j = this->pos(c);		
		
    // check (valid nucleotide && has child && child has probability table)
    if ( j >= 0 && N->child_[j] != NULL && N->child_[j]->has_table(frame) ) {
      stop = false;
      N = N->child_[j];
    }
    else break;
    i--;
  }
	
  char c = s[pos];
  int j = this->pos(c);
  return log(N->f_[frame](j));
}

double vom::VOM::log_prob(const std::string& s)
{
  double l = 0;
  for (int i = 0; i < sz(s); i++)
    l += log_prob(s, i);
  return l;
}

int vom::VOM::pos(char c)
{
  for (int j = 0; j < sz(alphabet); j++) 
    if (alphabet[j] == c) return j;
  return -1;
}

char vom::VOM::ch(int pos)
{
  return alphabet[pos];
}

void vom::VOM::grow_node(node* N, const std::vector<vom::sastring>& train_set)
{
  if (N == NULL) return;
  count_ ++;
	
  for (int i = 0; i < sz(alphabet); i++)
    if (N->child_[i] == NULL) try_grow(N, i, train_set);
	
  BOOST_FOREACH(node* child, N->child_)
    grow_node(child, train_set);
}

void vom::VOM::try_grow(node* N, int pos, const std::vector<vom::sastring>& train_set)
{
  char a = ch(pos);
  node* child = new node(a, alphabet.c_str(), N, periodic_);
  child->set_frequency(train_set);
  if (child->n_occurence_ < 2) {
    delete child;
    return;
  }

  /* 
  // Grow every periodic	
  double KL = 0;
  for (int pr = 0; pr < periodic_; pr++)
    KL += RAND::KLdivergence(child->f_[pr], N->f_[pr]);

  if (KL > eta_*periodic_) {
    //MSG("grow " << a << " " << N->str() << " KL=" << KL);
    N->child_[pos] = child;
  }
  else
    delete child;
  */
  bool ok = false;
  for (int pr = 0; pr < periodic_; pr++) 
    if (N->has_table(pr) && child->has_table(pr)) { // if N && child have this periodic table
      double KL = RAND::KLdivergence(child->f_[pr], N->f_[pr]);
      if (KL > eta_) ok = true;
      else child->f_[pr] = arma::vec(); // empty probability table
    }
    else child->f_[pr] = arma::vec(); // empty probability table
  // if some periodic table is informative then add the child
  if (ok)
    N->child_[pos] = child;
  else
    delete child;
}

void vom::VOM::prune_node(node** N, const std::vector<vom::sastring>& train_set)
{
  if (*N == NULL) return;
  for (int i = 0; i < sz(alphabet); i++) {
    node** child = &((*N)->child_[i]);
    prune_node(child, train_set);
  }
	
  /*
  bool no_child = true;
  for (int i = 0; i < sz(alphabet); i++)
    if ((*N)->child_[i] != NULL) no_child = false;
	
  if (!no_child) return;
  */
	
  node* parent = (*N)->parent_;
  if (parent != NULL) {
    /*
    double KL = 0;
    for (int pr = 0; pr < periodic_; pr++)
      KL += RAND::KLdivergence((*N)->f_[pr], parent->f_[pr]);
    if (KL < eta_*periodic_) {
      delete *N;
      *N = NULL;
      count_--;
    }
    */
    for (int pr = 0; pr < periodic_; pr++) {
      // if at least one child of N has this periodic table then we cannot prune
      bool ok = true;
      for (int i = 0; i < sz(alphabet); i++) 
	if ((*N)->child_[i] != NULL && (*N)->child_[i]->has_table(pr)) 
	  { ok = false; break; }
      if (!ok) continue;

      // otherwise, prune this periodic table if KL distance is small
      if ((*N)->has_table(pr) && parent->has_table(pr)) {
	double KL = RAND::KLdivergence((*N)->f_[pr], parent->f_[pr]);
	if (KL <= eta_) {
	  (*N)->f_[pr] = arma::vec();
	}
      }
    }

    // if all periodic tables are pruned then prune the node
    bool ok = false;
    for (int pr = 0; pr < periodic_; pr++) 
      if ((*N)->has_table(pr)) { ok = true; break; }
    if (!ok) {
      delete *N;
      *N = NULL;
      count_--;
    }
  }
}

const vom::FrequencyTable& vom::VOM::f_table(const std::string& s, int pos)
{
  node* N = root_;
  int i = pos-1;
	
  bool stop = false;
  while (i >= 0 && !stop) { // trace upstream for longest context
    char c = s[i];
    stop = true;
    int j = this->pos(c);		
		
    if (j >= 0 && N->child_[j] != NULL) {
      stop = false;
      N = N->child_[j];
    }
    else break;
    i--;
  }
  return N->f_[pos%periodic_];
}

void vom::VOM::simulate(int len, std::string& s)
{
  const int INIT_SIZE = 20;
  std::string tmp(len + INIT_SIZE, ' ');

  for (int i = 0; i < sz(tmp); i++) {
    const FrequencyTable& f = f_table(tmp, i);
    int j = RAND::random_idx(f);
    tmp.at(i) = alphabet[j];
  }

  s = tmp.substr(INIT_SIZE);// only take the sequence after INIT_SIZE characters
}

void vom::VOM::depth_node(node* N, int depth)
{
  if (N == NULL) return;
  count_++;
  if (depth_ < depth) depth_ = depth;
  for (int pr = 0; pr < periodic_; pr++) {
    if (N->has_table(pr)) {
      count_p_[pr]++;
      if (depth_p_[pr] < depth) depth_p_[pr] = depth;
    }
  }

  BOOST_FOREACH(node* child, N->child_) 
    depth_node(child, depth+1);
}
